#!/usr/bin/python
#coding:utf-8
'''
调参技巧：卷积神经网络-VGGNet实例，多分类器,tensorboard功能，web查看网页内容,引入数据增强，计算图增强图片,加深卷机层，6个卷积层扩展为9个卷机层，
'''
import tensorflow as tf 
import os 
import _pickle as cPickle 
import numpy as np 

CIFAR_DIR="./cifar-10-batches-py"

# tensorboard 实现不周
# 1、指定面板图上显示的变量
# 2、训练过程中将变量计算出来，输出到文件中
# 3、文件解析  ./tensorboard  --logdir=dir

def load_data(filename):
    '''
    从文件读取数据
    '''
    with open(filename,'rb') as f:
        data=cPickle.load(f,encoding='bytes')
        return data[b'data'],data[b'labels']

class CifarData:
    '''
    数据处理类
    '''
    def __init__(self,fileList,need_shuffle): #need_shuffle 是否需要shuffle
        all_data=[]
        all_labelList = []
        for filename in fileList:
            data,lableList= load_data(filename)
            
            #获取全量十个类数据
            all_data.append(data)
            all_labelList.append(lableList)

            #获取两类数据
            # for item,label in zip(data,lableList):
            #     if label in [0,1]:
            #         all_data.append(item)
            #         all_labelList.append(label)
        
        self._data= np.vstack(all_data)  #纵行合并为矩阵
        self._data = self._data  
        self._labels= np.hstack(all_labelList) #横向合并为一维向量
        self._num_example=self._data.shape[0]
        self._need_shuffle = need_shuffle
        self._indicator = 0

        if self._need_shuffle:
            self._shuffle_data()
    
    def _shuffle_data(self):
        '''
        打乱数据
        '''
        p = np.random.permutation(self._num_example)  #实现数据混排
        self._data=self._data[p]
        self._labels=self._labels[p]
    
    def nex_batche(self,batch_size):
        '''
        每次返回batch_size个样本
        数据获取一遍后，shuffle一遍，重新获取
        '''
        end_indicator= self._indicator + batch_size
        if end_indicator > self._num_example:
            if self._need_shuffle:
                self._shuffle_data()
                self._indicator=0
                end_indicator = batch_size
            else:
                raise Exception('have no more examples')
        
        if end_indicator > self._num_example:
            raise Exception("batch size is larger than all example")

        batch_data=self._data[self._indicator:end_indicator]
        batch_labels = self._labels[self._indicator: end_indicator]
        self._indicator=end_indicator

        return batch_data,batch_labels

train_filenames=[os.path.join(CIFAR_DIR,'data_batch_%d'% i) for i in range(1,6)]
test_filenames =[os.path.join(CIFAR_DIR,'test_batch')]

train_data=CifarData(train_filenames,True)
test_data=CifarData(test_filenames,False)

batch_size=20
#构建计算图
x= tf.placeholder(tf.float32,[batch_size,3072]) #placeholder(占位符
y= tf.placeholder(tf.int64,[batch_size]) #None 样本数不确定

#数据转换
x_image=tf.reshape(x,[-1,3,32,32])   # 格式转换
x_image=tf.transpose(x_image,perm=[0,2,3,1])  #通道转换，编程正确顺序
x_image_arr= tf.split(x_image,num_or_size_splits=batch_size,axis=0)

result_x_image_arr=[]
for x_single_image in x_image_arr:
    # x_single_image[1,32,32,3]
    x_single_image =  tf.reshape(x_single_image,[32,32,3])

    ######图像增强##########
    data_aug_1 = tf.image.random_flip_left_right(x_single_image) #随机左右翻转
    data_aug_2 = tf.image.random_brightness(data_aug_1,max_delta=63) #随机调整光照强度,max_delta光照强度参数
    data_aug_3 = tf.image.random_contrast(data_aug_2,lower=0.2,upper=1.8) #随机对比度,调整区间最大1.8，最小0.2

    x_single_image = tf.reshape(data_aug_3,[1,32,32,3])
    result_x_image_arr.append(x_single_image)

result_x_images = tf.concat(result_x_image_arr,axis=0) #数据合并
normal_result_x_images =result_x_images /127.5 -1


#第一层 卷积层 conv1：神经元图 feature_map 输出图像
conv1_1= tf.layers.conv2d(normal_result_x_images,
                        32,  #output channel number 输出通道数
                        (3,3), #kernel size卷积核大小
                        padding='same', #padding类型，same（做padding）、valid（不做padding）
                        activation = tf.nn.relu, #激活函数
                        name= 'conv1_1')

conv1_2= tf.layers.conv2d(conv1_1,
                        32,  #output channel number 输出通道数
                        (3,3), #kernel size卷积核大小
                        padding='same', #padding类型，same（做padding）、valid（不做padding）
                        activation = tf.nn.relu, #激活函数
                        name= 'conv1_2')
###新增卷积层
conv1_3= tf.layers.conv2d(conv1_2,
                        32,  #output channel number 输出通道数
                        (3,3), #kernel size卷积核大小
                        padding='same', #padding类型，same（做padding）、valid（不做padding）
                        activation = tf.nn.relu, #激活函数
                        name= 'conv1_3')

#第一层 池化层 大小减小一半 16*16
pooling1 =tf.layers.max_pooling2d(conv1_3,
                        (2,2), #kernel size 内核大小
                        (2,2), #stride 步长
                        name='pool1')


#第二层 卷积层 conv1：神经元图 feature_map 输出图像
conv2_1 = tf.layers.conv2d(pooling1,
                        32,  #output channel number 输出通道数
                        (3,3), #kernel size卷积核大小
                        padding='same', #padding类型，same（做padding）、valid（不做padding）
                        activation = tf.nn.relu, #激活函数
                        name= 'conv2_1')

conv2_2 = tf.layers.conv2d(conv2_1,
                        32,  #output channel number 输出通道数
                        (3,3), #kernel size卷积核大小
                        padding='same', #padding类型，same（做padding）、valid（不做padding）
                        activation = tf.nn.relu, #激活函数
                        name= 'conv2_2')

conv2_3 = tf.layers.conv2d(conv2_2,
                        32,  #output channel number 输出通道数
                        (3,3), #kernel size卷积核大小
                        padding='same', #padding类型，same（做padding）、valid（不做padding）
                        activation = tf.nn.relu, #激活函数
                        name= 'conv2_3')

#第二层 池化层 大小减小一半 8*8
pooling2 =tf.layers.max_pooling2d(conv2_3,
                        (2,2), #kernel size 内核大小
                        (2,2), #stride 步长
                        name='pool2')

#第二层 卷积层 conv1：神经元图 feature_map 输出图像
conv3_1 = tf.layers.conv2d(pooling2,
                        32,  #output channel number 输出通道数
                        (3,3), #kernel size卷积核大小
                        padding='same', #padding类型，same（做padding）、valid（不做padding）
                        activation = tf.nn.relu, #激活函数
                        name= 'conv3_1')

conv3_2 = tf.layers.conv2d(conv3_1,
                        32,  #output channel number 输出通道数
                        (3,3), #kernel size卷积核大小
                        padding='same', #padding类型，same（做padding）、valid（不做padding）
                        activation = tf.nn.relu, #激活函数
                        name= 'conv3_2')

conv3_3 = tf.layers.conv2d(conv3_2,
                        32,  #output channel number 输出通道数
                        (3,3), #kernel size卷积核大小
                        padding='same', #padding类型，same（做padding）、valid（不做padding）
                        activation = tf.nn.relu, #激活函数
                        name= 'conv3_3')

#第二层 池化层 大小减小一半 4*4 32个
pooling3 =tf.layers.max_pooling2d(conv3_3,
                        (2,2), #kernel size 内核大小
                        (2,2), #stride 步长
                        name='pool3')

# 展平 [None,4* 4 * 32]
flatten = tf.layers.flatten(pooling3)

#全连接层
y_ = tf.layers.dense(flatten,10)

###########平方差损失函数###########
#针对当前数据集，该损失函验证集成功率40%
#api: e^x / sum(e^x)
# p_y= tf.nn.softmax(y_)
# #5 --> [0000100000]
# y_one_hot= tf.one_hot(y,10,dtype=tf.float32)
# loss=tf.reduce_mean(tf.square(y_one_hot - p_y))  #reduce_mean 求均值  square：求平方，结果为真实值和预测值差的平方

###########交叉熵损失函数##########
#针对当前数据集，损失函数验证集成功率30%
loss=tf.losses.sparse_softmax_cross_entropy(labels=y,logits=y_)

'''
#sigmoid 
#[None,10]
p_y_1 = tf.nn.sigmoid(y_)  #传入sigmoid得到概率值
#[None,1]
y_reshaped=tf.reshape(y,(-1,1))  #数据类型 int64
y_reshaped_float = tf.cast(y_reshaped,tf.float32)  #类型转换 int64 转 float32

#平方差计算损失函数
loss=tf.reduce_mean(tf.square(y_reshaped_float - p_y_1))  #reduce_mean 求均值  square：求平方，结果为真实值和预测值差的平方
'''

#准确率
predict =tf.argmax(y_,1) 

#[1,0,1,1...]
correct_prediction= tf.equal(predict,y) #得到正确预测值
accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float64)) #得到准确率

#梯度下降方法
with tf.name_scope('train_op'):
    train_op = tf.train.AdamOptimizer(1e-3).minimize(loss)  #1e-3：10的负三次方

##################指定tensorborad参数#############

def variable_summary(var,name):
    '''
    tensorboard封装方法，为一个变量统计量指定summary
    '''
    with tf.name_scope(name):
        mean= tf.reduce_mean(var) #均值
        with tf.name_scope('stddev'): #平方差
            stddev= tf.sqrt(tf.reduce_mean(tf.square(var-mean)))
        tf.summary.scalar('mean',mean)
        tf.summary.scalar('stddev',stddev)
        tf.summary.scalar('min',tf.reduce_min(var))
        tf.summary.scalar('max',tf.reduce_max(var))
        tf.summary.histogram('histogram',var)  #指定直方图

#为各个输出层增加summary
with tf.name_scope('summary'):
    variable_summary(conv1_1,'conv1_1')
    variable_summary(conv1_2,'conv1_2')
    variable_summary(conv2_1,'conv2_1')
    variable_summary(conv2_2,'conv2_2')
    variable_summary(conv3_1,'conv3_1')
    variable_summary(conv3_2,'conv3_2')

loss_summary= tf.summary.scalar('loss',loss)
#summary 聚合每一步参数
accuracy_summary = tf.summary.scalar('accuracy',accuracy)

input_summary = tf.summary.image('inputs_images',result_x_images)

merged_summary = tf.summary.merge_all() #方式一，所有summary对象，都加入
merged_summary_test = tf.summary.merge([loss_summary,accuracy_summary]) #方式二，选择性summary

def checkDir(path):
    '''
    判断目录并创建
    '''
    if not  os.path.exists(path):
        os.mkdir(path)

#输出到文件
LOG_DIR = './'
run_label = 'run_vgg_tensorboard'
run_dir = os.path.join(LOG_DIR,run_label)
checkDir(run_dir)
train_log_dir = os.path.join(run_dir,'train')
test_log_dir = os.path.join(run_dir,'test')
checkDir(train_log_dir)
checkDir(test_log_dir)

#创建绘画
init = tf.global_variables_initializer()

train_steps=100000
test_steps=100

#指定summary多少步输出一次
output_summary_every_steps= 100

#train 100k: 78.04%，数据增强后：82.6%，扩展卷机层后成功率 83.4%
with tf.Session() as sess:
    sess.run(init)
    ###tensorboard
    train_writer= tf.summary.FileWriter(train_log_dir,sess.graph) #指定计算图
    test_wirter = tf.summary.FileWriter(test_log_dir) #不指定计算图

    #取出并存储测试数据集
    fixed_test_batch_data,fixed_test_batch_labels= test_data.nex_batche(batch_size)

    for i in range(train_steps):
        batch_data,batch_labels=train_data.nex_batche(batch_size)

        #tensorboard
        eval_ops = [loss,accuracy,train_op]
        should_output_summary = ((i+1)% output_summary_every_steps==0)
        if should_output_summary:
            eval_ops.append(merged_summary)

        eval_ops_results =sess.run(
                    eval_ops,
                    # [loss,accuracy,train_op],
                    feed_dict={x:batch_data, y:batch_labels})  #加train_op 表示训练，不加表示计算，feed_dict:输入数据
        
        loss_val,acc_val =eval_ops_results[0:2]

        #tensorboard
        if should_output_summary:
            #训练输出到文件
            train_summary_str= eval_ops_results[3]
            train_writer.add_summary(train_summary_str,i+1)
            #测试输出到文件
            test_summary_str= sess.run([merged_summary_test],
                                        feed_dict={
                                            x:fixed_test_batch_data,
                                            y:fixed_test_batch_labels
                                        })[0]
            
            test_wirter.add_summary(test_summary_str,i+1)

        if i % 100 ==0:
            print("[Train] Step: %d, loss: %4.5f, acc: %4.5f" %(i,loss_val,acc_val))
        
        #每5000次test数据集评测
        if (i+1) % 1000 ==0:
            test_data=CifarData(test_filenames,False)
            all_test_acc_val=[]
            for j in range(test_steps):
                test_batch_data,test_batch_labels=test_data.nex_batche(batch_size)
                test_acc_val=sess.run(
                    [accuracy],
                    feed_dict={
                        x:test_batch_data,
                        y:test_batch_labels
                    })
                all_test_acc_val.append(test_acc_val)
            
            test_acc=np.mean(all_test_acc_val)
            print('[Test ] Step: %d, acc: %4.5f' % (i+1,test_acc))